#include "StdAfx.h"
#include "ViewProcessesBySnapshot.h"
#include <windows.h>
#include <tlhelp32.h>
#include <tchar.h>

CViewProcessesBySnapshot::CViewProcessesBySnapshot(void)
{
	m_eLogType = LogType_File;
	TCHAR* pFileName = _T("ProcessInfos.dat");
	if(!m_fLogFile.Open(pFileName, CFile::modeCreate | CFile::modeWrite 
		| CFile::typeText)) 
	{
		TRACE(_T("Unable to open file\n"));
	}
}


CViewProcessesBySnapshot::~CViewProcessesBySnapshot(void)
{
	m_fLogFile.Close();
}

//  Forward declarations:
//BOOL GetProcessList( );
//BOOL ListProcessModules( DWORD dwPID );
//BOOL ListProcessThreads( DWORD dwOwnerPID );
//void printError( TCHAR* msg );

int CViewProcessesBySnapshot::ListProcesses( void )
{
	GetProcessList( );
	return 0;
}

BOOL CViewProcessesBySnapshot::GetProcessList( )
{
	HANDLE hProcessSnap;
	HANDLE hProcess;
	PROCESSENTRY32 pe32;
	DWORD dwPriorityClass;

	// Take a snapshot of all processes in the system.
	hProcessSnap = CreateToolhelp32Snapshot( TH32CS_SNAPPROCESS, 0 );
	if( hProcessSnap == INVALID_HANDLE_VALUE )
	{
		PrintError( TEXT("CreateToolhelp32Snapshot (of processes)") );
		return( FALSE );
	}

	// Set the size of the structure before using it.
	pe32.dwSize = sizeof( PROCESSENTRY32 );

	// Retrieve information about the first process,
	// and exit if unsuccessful
	if( !Process32First( hProcessSnap, &pe32 ) )
	{
		PrintError( TEXT("Process32First") ); // show cause of failure
		CloseHandle( hProcessSnap );          // clean the snapshot object
		return( FALSE );
	}

	CString strTmp;
	// Now walk the snapshot of processes, and
	// display information about each process in turn
	do
	{
		PrintLog( TEXT("\n\n=====================================================" ), m_eLogType);
		strTmp.Format(_T("\nPROCESS NAME:  %s"), pe32.szExeFile );
		PrintLog(strTmp, m_eLogType);
		PrintLog( TEXT("\n-------------------------------------------------------" ), m_eLogType);

		// Retrieve the priority class.
		dwPriorityClass = 0;
		hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, pe32.th32ProcessID );
		if( hProcess == NULL )
			PrintError( TEXT("OpenProcess") );
		else
		{
			dwPriorityClass = GetPriorityClass( hProcess );
			if( !dwPriorityClass )
				PrintError( TEXT("GetPriorityClass") );
			CloseHandle( hProcess );
		}

		//_tprintf( TEXT("\n  Process ID        = 0x%08X"), pe32.th32ProcessID );
		//_tprintf( TEXT("\n  Thread count      = %d"),   pe32.cntThreads );
		//_tprintf( TEXT("\n  Parent process ID = 0x%08X"), pe32.th32ParentProcessID );
		//_tprintf( TEXT("\n  Priority base     = %d"), pe32.pcPriClassBase );

		strTmp.Format(_T("\n Process ID        = 0x%08X"), pe32.th32ProcessID );
		PrintLog(strTmp, m_eLogType);
		strTmp.Format(_T("\n  Thread count      = %d"),   pe32.cntThreads );
		PrintLog(strTmp, m_eLogType);
		strTmp.Format(_T("\n  Parent process ID = 0x%08X"), pe32.th32ParentProcessID );
		PrintLog(strTmp, m_eLogType);
		strTmp.Format(_T("\n  Priority base     = %d"), pe32.pcPriClassBase );
		PrintLog(strTmp, m_eLogType);

		if( dwPriorityClass )
			_tprintf( TEXT("\n  Priority class    = %d"), dwPriorityClass );

		// List the modules and threads associated with this process
		ListProcessModules( pe32.th32ProcessID );
		ListProcessThreads( pe32.th32ProcessID );

	} while( Process32Next( hProcessSnap, &pe32 ) );

	CloseHandle( hProcessSnap );
	return( TRUE );
}


BOOL CViewProcessesBySnapshot::ListProcessModules( DWORD dwPID )
{
	HANDLE hModuleSnap = INVALID_HANDLE_VALUE;
	MODULEENTRY32 me32;

	// Take a snapshot of all modules in the specified process.
	hModuleSnap = CreateToolhelp32Snapshot( TH32CS_SNAPMODULE, dwPID );
	if( hModuleSnap == INVALID_HANDLE_VALUE )
	{
		PrintError( TEXT("CreateToolhelp32Snapshot (of modules)") );
		return( FALSE );
	}

	// Set the size of the structure before using it.
	me32.dwSize = sizeof( MODULEENTRY32 );

	// Retrieve information about the first module,
	// and exit if unsuccessful
	if( !Module32First( hModuleSnap, &me32 ) )
	{
		PrintError( TEXT("Module32First") );  // show cause of failure
		CloseHandle( hModuleSnap );           // clean the snapshot object
		return( FALSE );
	}
	CString strTmp;
	// Now walk the module list of the process,
	// and display information about each module
	do
	{
		//_tprintf( TEXT("\n\n     MODULE NAME:     %s"),   me32.szModule );
		//_tprintf( TEXT("\n     Executable     = %s"),     me32.szExePath );
		//_tprintf( TEXT("\n     Process ID     = 0x%08X"),         me32.th32ProcessID );
		//_tprintf( TEXT("\n     Ref count (g)  = 0x%04X"),     me32.GlblcntUsage );
		//_tprintf( TEXT("\n     Ref count (p)  = 0x%04X"),     me32.ProccntUsage );
		//_tprintf( TEXT("\n     Base address   = 0x%08X"), (DWORD) me32.modBaseAddr );
		//_tprintf( TEXT("\n     Base size      = %d"),             me32.modBaseSize );

		strTmp.Format(_T("\n\n     MODULE NAME:     %s"),   me32.szModule );
		PrintLog(strTmp, m_eLogType);
		strTmp.Format(_T("\n     Executable     = %s"),     me32.szExePath );
		PrintLog(strTmp, m_eLogType);
		strTmp.Format(_T("\n     Process ID     = 0x%08X"),         me32.th32ProcessID );
		PrintLog(strTmp, m_eLogType);
		strTmp.Format(_T("\n     Ref count (g)  = 0x%04X"),     me32.GlblcntUsage );
		PrintLog(strTmp, m_eLogType);
		strTmp.Format(_T("\n     Ref count (p)  = 0x%04X"),     me32.ProccntUsage );
		PrintLog(strTmp, m_eLogType);
		strTmp.Format(_T("\n     Base address   = 0x%08X"), (DWORD) me32.modBaseAddr );
		PrintLog(strTmp, m_eLogType);
		strTmp.Format(_T("\n     Base size      = %d"),             me32.modBaseSize );
		PrintLog(strTmp, m_eLogType);

	} while( Module32Next( hModuleSnap, &me32 ) );

	CloseHandle( hModuleSnap );
	return( TRUE );
}

BOOL CViewProcessesBySnapshot::ListProcessThreads( DWORD dwOwnerPID ) 
{ 
	HANDLE hThreadSnap = INVALID_HANDLE_VALUE; 
	THREADENTRY32 te32; 

	// Take a snapshot of all running threads  
	hThreadSnap = CreateToolhelp32Snapshot( TH32CS_SNAPTHREAD, 0 ); 
	if( hThreadSnap == INVALID_HANDLE_VALUE ) 
		return( FALSE ); 

	// Fill in the size of the structure before using it. 
	te32.dwSize = sizeof(THREADENTRY32); 

	// Retrieve information about the first thread,
	// and exit if unsuccessful
	if( !Thread32First( hThreadSnap, &te32 ) ) 
	{
		PrintError( TEXT("Thread32First") ); // show cause of failure
		CloseHandle( hThreadSnap );          // clean the snapshot object
		return( FALSE );
	}

	CString strTmp;
	// Now walk the thread list of the system,
	// and display information about each thread
	// associated with the specified process
	do 
	{ 
		if( te32.th32OwnerProcessID == dwOwnerPID )
		{
			//_tprintf( TEXT("\n\n     THREAD ID      = 0x%08X"), te32.th32ThreadID ); 
			//_tprintf( TEXT("\n     Base priority  = %d"), te32.tpBasePri ); 
			//_tprintf( TEXT("\n     Delta priority = %d"), te32.tpDeltaPri ); 
			//_tprintf( TEXT("\n"));

			strTmp.Format(_T("\n\n     THREAD ID      = 0x%08X"), te32.th32ThreadID ); 
			PrintLog(strTmp, m_eLogType);
			strTmp.Format(_T("\n     Base priority  = %d"), te32.tpBasePri ); 
			PrintLog(strTmp, m_eLogType);
			strTmp.Format(_T("\n     Delta priority = %d"), te32.tpDeltaPri ); 
			PrintLog(strTmp, m_eLogType);
			strTmp.Format(_T("\n") );
			PrintLog(strTmp, m_eLogType);
		}
	} while( Thread32Next(hThreadSnap, &te32 ) ); 

	CloseHandle( hThreadSnap );
	return( TRUE );
}

void CViewProcessesBySnapshot::PrintError( TCHAR* msg )
{
	DWORD eNum;
	TCHAR sysMsg[256];
	TCHAR* p;

	eNum = GetLastError( );
	FormatMessage( FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL, eNum,
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), // Default language
		sysMsg, 256, NULL );

	// Trim the end of the line and terminate it with a null
	p = sysMsg;
	while( ( *p > 31 ) || ( *p == 9 ) )
		++p;
	do { *p-- = 0; } while( ( p >= sysMsg ) &&
		( ( *p == '.' ) || ( *p < 33 ) ) );

	// Display the message
	_tprintf( TEXT("\n  WARNING: %s failed with error %d (%s)"), msg, eNum, sysMsg );
}

void CViewProcessesBySnapshot::PrintLog( LPCTSTR p_strLog, E_LogType p_eLogType /*= LogType_Console*/ )
{
	switch(p_eLogType)
	{
	case LogType_Console:
		{
			_tprintf( p_strLog);
			break;
		}
	case LogType_File:
		{
			m_fLogFile.WriteString(p_strLog);
			break;
		}
	default:
		{
			break;
		}
	}
}
